import numpy as np
import torchvision.datasets as dsets
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import seaborn as sns
import random

def visualize_xgboost_predictions(X_test, y_test, y_test_pred, num_images=10):
    """可视化XGBoost预测结果，显示原始图像和识别的数字"""
    # 将展平的图像重新整形为28x28
    images = X_test[:num_images].reshape(num_images, 28, 28)
    
    # 创建子图
    fig, axes = plt.subplots(2, 5, figsize=(15, 8))
    axes = axes.ravel()
    
    for i in range(num_images):
        # 显示图像
        axes[i].imshow(images[i], cmap='gray')
        
        # 设置标题显示真实标签和预测标签
        color = 'green' if y_test[i] == y_test_pred[i] else 'red'
        axes[i].set_title(f'True: {y_test[i]}, Pred: {y_test_pred[i]}', color=color, fontsize=12)
        axes[i].axis('off')
    
    plt.suptitle('XGBoost MNIST Digit Recognition Results', fontsize=16)
    plt.tight_layout()
    plt.show()

def plot_confusion_matrix(y_true, y_pred):
    """绘制混淆矩阵"""
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.title('Confusion Matrix for XGBoost MNIST Classification')
    plt.xlabel('Predicted Label')
    plt.ylabel('True Label')
    plt.show()

def main():
    # 设置随机种子以确保结果可重现
    random_seed = np.random.randint(0, 10000)
    np.random.seed(random_seed)
    random.seed(random_seed)
    
    # 设置批次大小
    batch_size = 1000  # 增加批次大小以加快处理速度
    
    # 加载MNIST数据集
    print("Loading MNIST dataset...")
    mnist_train_dataset = dsets.MNIST(root="dataset/mnist", train=True, download=True)
    mnist_test_dataset = dsets.MNIST(root="dataset/mnist", train=False, download=True)
    
    # 创建数据加载器
    train_loader = DataLoader(dataset=mnist_train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(dataset=mnist_test_dataset, batch_size=batch_size, shuffle=True)  # 改为True以随机选择测试数据
    
    # 设置随机选择的数据量
    train_sample_size = 10000  # 从60000个训练样本中随机选择10000个
    test_sample_size = 2000    # 从10000个测试样本中随机选择2000个
    
    # 准备训练数据（随机选择）
    print(f"Preparing training data (randomly selecting {train_sample_size} samples)...")
    train_indices = np.random.choice(len(train_loader.dataset), train_sample_size, replace=False)
    X_train = train_loader.dataset.data[train_indices].numpy()
    X_train = X_train.reshape(X_train.shape[0], 28 * 28)  # 展平图像为784维特征向量
    y_train = train_loader.dataset.targets[train_indices].numpy()
    
    # 准备测试数据（随机选择）
    print(f"Preparing test data (randomly selecting {test_sample_size} samples)...")
    test_indices = np.random.choice(len(test_loader.dataset), test_sample_size, replace=False)
    X_test = test_loader.dataset.data[test_indices].numpy()
    X_test = X_test.reshape(X_test.shape[0], 28 * 28)  # 展平图像为784维特征向量
    y_test = test_loader.dataset.targets[test_indices].numpy()
    
    # 创建XGBoost分类器
    print("Creating XGBoost classifier...")
    xgb_model = XGBClassifier(
        n_estimators=10,             # 减少树的数量以提高训练速度
        max_depth=6,                # 树的最大深度
        learning_rate=0.1,          # 学习率
        subsample=0.8,              # 训练样本采样比例
        colsample_bytree=0.8,       # 特征采样比例
        random_state=random_seed,   # 使用与数据选择相同的随机种子
        verbosity=1,                # 显示训练过程信息
        n_jobs=-1                   # 使用所有CPU核心并行训练
    )
    
    # 训练模型
    print("Training XGBoost model...")
    print("Note: Training may take 1-3 minutes depending on your hardware.")
    print("Training progress (showing every 10 iterations):")
    xgb_model.fit(X_train, y_train, verbose=10)  # 每10轮显示一次进度
    
    # 进行预测
    print("Making predictions...")
    y_test_pred = xgb_model.predict(X_test)
    
    # 计算准确率
    accuracy = accuracy_score(y_test, y_test_pred)
    print(f'Accuracy: {accuracy:.4f}')
    
    # 打印详细分类报告
    print("\nClassification Report:")
    print(classification_report(y_test, y_test_pred))
    
    # 可视化预测结果（显示前10个测试样本）
    print("Visualizing predictions...")
    visualize_xgboost_predictions(X_test, y_test, y_test_pred, num_images=10)
    
    # 绘制混淆矩阵
    print("Plotting confusion matrix...")
    plot_confusion_matrix(y_test, y_test_pred)

if __name__ == "__main__":
    main()